package com.example.springai.config;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.milvus.MilvusVectorStore;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.context.annotation.Bean;

/**
 * 聊天配置
 *
 * @author ZWx
 */
@Slf4j
@SpringBootConfiguration
@RequiredArgsConstructor
public class ChatConfig {

    private final ChatModel chatModel;
    private final ChatMemoryRepository chatMemoryRepository;
    private final ToolCallbackProvider tools;
    private final VectorStore vectorStore;

    /**
     * 创建 chatClient 聊天客户端
     *
     * @return 聊天客户端
     */
    @Bean
    public ChatClient chatClient() {
        log.info("======================== Tools ==========================");
        for (ToolCallback toolCallback : tools.getToolCallbacks()) {
            log.info("Tool Definition: {}", toolCallback.getToolDefinition());
            log.info("Tool Metadata: {}", toolCallback.getToolMetadata());
        }
        log.info("======================== Tools ==========================");

        // 创建 chatMemory 聊天记忆
        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(chatMemoryRepository)
                .maxMessages(5)
                .build();

        // 创建 vectorStoreDocumentRetriever 向量存储检索器
        VectorStoreDocumentRetriever vectorStoreDocumentRetriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .topK(5)
                .build();

        // 创建 queryAugmenter 查询增强
        ContextualQueryAugmenter contextualQueryAugmenter = ContextualQueryAugmenter.builder()
                .allowEmptyContext(true)
                .build();

        // 创建 rag 检索增强
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(vectorStoreDocumentRetriever)
                .queryAugmenter(contextualQueryAugmenter)
                .build();

        return ChatClient.builder(chatModel)
                .defaultSystem("你是一个学识渊博的AI助手，使用中文回答问题。")
                // 注册Advisor
                // 默认日志
                .defaultAdvisors(new SimpleLoggerAdvisor())
                // 聊天记忆
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build())
                // rag
                .defaultAdvisors(retrievalAugmentationAdvisor)
                // tools
                .defaultToolCallbacks(tools.getToolCallbacks())
                .build();
    }
}
